import pickle
from easy_io import write_pkl_file
import numpy as np


def bytes_to_str(b):
    if isinstance(b, bytes):
        return b.decode()
    elif isinstance(b, np.str_):
        return str(b)
    else:
        raise ValueError(b)


def convert(candidates):
    candidates = [dict(source=bytes_to_str(c[b'source']), scanid=bytes_to_str(c[b'scan_id'] if b'scan_id' in c else c[b'scanid']), bbox=np.asarray(c[b'bbox'], 'int'))
                  for c in candidates]
    return candidates


def check(candidates):
    assert all(isinstance(c['bbox'], np.ndarray) and c['bbox'].shape==(6,) and c['bbox'].dtype==np.int for c in candidates)
    assert all(np.all(c['bbox'] >= 0) & np.all(c['bbox'][-3:] > c['bbox'][:3]) for c in candidates)
    assert all(isinstance(c['source'], str) and isinstance(c['scanid'], str) for c in candidates)


def main(infile, outfile, loader=pickle):
    with open(infile, 'rb') as f:
        candidates = loader.load(f, encoding='bytes')
    candidates = convert(candidates)
    print('nb_candidates', len(candidates))
    candidates = [c for c in candidates if c['scanid'] != 'b8bb02d229361a623a4dc57aa0e5c485']
    check(candidates)
    print('nb_candidates', len(candidates))
    write_pkl_file(outfile, candidates)


if __name__ == '__main__':
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_list_spie.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v1.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_list_kaggle_testset.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v1.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_list_kaggle_beni.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v1.pkl',
    # )

    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_5fold_with_lungmask.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v8.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_kaggle_beni_with_lungmask.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v3.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_kaggle_testset_with_lungmask.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v3.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_spie_with_lungmask.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v3.pkl',
    # )

    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_5fold_with_lungmask_0328.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v10.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_kaggle_beni_with_lungmask_0328.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v5.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_kaggle_testset_with_lungmask_0328.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v5.pkl',
    # )
    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_spie_with_lungmask_0328.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v5.pkl',
    # )

    # main(
    #     infile='/data_4t/Kaggle/candidates/candidates_5fold_with_lungmask_0330_thresh_0.1_recall_704_710_fp_68.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v12.pkl'
    # )

    # main(
    #     infile='/data-disk/Kaggle/candidates/candidates_kaggle_beni_with_lungmask_0330.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v7.pkl',
    # )
    #
    # main(
    #     infile='/data-disk/Kaggle/candidates/candidates_kaggle_testset_with_lungmask_0330.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v7.pkl',
    # )
    #
    # main(
    #     infile='/data-disk/Kaggle/candidates/candidates_spie_with_lungmask_0330.pkl',
    #     outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v7.pkl',
    # )

    # main(
    #     infile='/data-disk/liax/candis_5fold_after_nms_lungmask_60000.npy',
    #     outfile='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v14.pkl',
    #     loader=np,
    # )

    main(
        infile='/data-disk/liax/candis_kaggle_beni_after_nms_lungmask_60000.npy',
        outfile='/ssd_1t/huzq/kaggle_data/kaggle_beni_candidates_v9.pkl',
        loader=np,
    )

    main(
        infile='/data-disk/liax/candis_kaggle_test_after_nms_lungmask_60000.npy',
        outfile='/ssd_1t/huzq/kaggle_data/kaggle_testset_candidates_v9.pkl',
        loader=np,
    )

    main(
        infile='/data-disk/liax/candis_spie_after_nms_lungmask_60000.npy',
        outfile='/ssd_1t/huzq/kaggle_data/spie_candidates_v9.pkl',
        loader=np,
    )
